import random
import math

def distance(p1, p2):
    return math.sqrt(sum((p1[i] - p2[i])**2 for i in range(len(p1))))

def assign_cluster(x, centers):
    distances = [distance(x, center) for center in centers]
    return distances.index(min(distances))


def Kmeans(data, k, epsilon=1e-4, iteration=100):
    centers = random.sample(data, k)
    
    for it in range(iteration):
        print(f"Iteration {it+1}")

        clusters = [[] for _ in range(k)]

        for point in data:
            idx = assign_cluster(point, centers)
            clusters[idx].append(point)

        new_centers = []
        shift = 0

        for i in range(k):
            if len(clusters[i]) == 0:
                new_centers.append(random.choice(data))
                continue

            dim = len(clusters[i][0])
            avg = [0] * dim
            for p in clusters[i]:
                for d in range(dim):
                    avg[d] += p[d]
            avg = [x / len(clusters[i]) for x in avg]

            shift += distance(centers[i], avg)
            new_centers.append(avg)

        print("Center shift:", shift)

        if shift < epsilon:
            print("Converged.")
            break

        centers = new_centers

    return centers, clusters
